#include "gepfit.h"

#define FSET_MAX_ARGS 2
#define SELECTION_RANGE 20

int main() {
	init_CL();
	load_CL("./gepfit.cl");
	loadSets();

	// INITIALIZE AND VERIFY POPULATION CHROMOSOMES
	printf("\n");
	population P = initPop(40, 50, 3);
//	popDump(&P);
	printf("\n");

	for (int gen = 0; gen < 100; gen++) {
		// MUTATION TEST
		recombine(&P);
		mutate(&P);
		//transpose(&P);

		// GET SUM FITNESS FOR EACH INDIVIDUAL
		computeFitness(&P);

		// VERIFY ASSEMBLY
		/*readBuffer(&P.B_assembly);
		int k = 0;
		for (int i = 0; i < P.count; i++) {
			for (int j = 0; j < (P.length+P.genes) * 4; j++) {
				printf("%i", ((unsigned char*)P.B_assembly.data)[i*((P.length+P.genes)*4) + j]);
				if (k%4 == 3)
					printf("; ");
				else	printf(".");
				k++;
			}
			printf("\n");
		}/**/

		// SELECT NEW POPULATION BY ELITIST ROULETTE REPLICATION
		replicate(&P);

	}
	
	// VERIFY NEW POPULATION
	printf("\n");
//	popDump(&P);
	printf("\n");

	// VERIFY CHAMPION
	printf("\n");
	rankSurvivors(&P);
	printf("\n");

	// CLEANUP
	deletePop(&P);
	kill_CL();

	printf("Exiting...\n");
	return EXIT_SUCCESS;
}

void loadSets() {
	FILE* T;
	FILE* F;
	FILE* D;
	
	T = fopen("./tset.txt", "r");
	F = fopen("./fset.txt", "r");
	D = fopen("./training.txt", "r");
	if (T == NULL) {
		printf("FATAL ERROR: No terminal symbols found!\n./tset.txt missing...\n");
		exit(EXIT_FAILURE);
	}
	if (F == NULL) {
		printf("FATAL ERROR: No function symbols found!\n./fset.txt missing...\n");
		exit(EXIT_FAILURE);
	}
	if (D == NULL) {
		printf("FATAL ERROR: No training data found!\n./training.txt missing...\n");
		exit(EXIT_FAILURE);
	}

	setF.symb = (char*) malloc (sizeof(char) * 256);

	setF.Fsize = fread(setF.symb, sizeof(char), 256, F) - 1;
	setF.Tsize = fread(setF.symb + setF.Fsize, sizeof(char), 256, T) - 1;
	*(setF.symb + setF.Fsize + setF.Tsize) = '\0';

	fclose(T);
	fclose(F);

	trainData.args = setF.Tsize;
	char read;
	int filesize =0, lines = 0;
	do {
		read = fgetc(D);
		if (read == '\n')
			lines++;
		filesize++;
	} while (read != EOF);
	rewind(D);
	trainData.size = lines;
	if (lines*(trainData.args + 1)*sizeof(float) > filesize) {
		printf("FATAL ERROR: Mismatch between "
		"training data and terminal symbols?\n");
		exit(EXIT_FAILURE);
	}
	trainData.B_data = initBuffer (sizeof(float) * (trainData.args+1) * trainData.size, INPUT_BUFFER);
	trainData.data = (float*) trainData.B_data.data;
	for (int i = 0; i < trainData.size; i++) {
		for (int j = 0; j < trainData.args; j++) {
			fscanf(D, "%f,", trainData.data + i*(trainData.args+1) + j);
		}
		fscanf(D, "%f\n", trainData.data + i*(trainData.args+1)+trainData.args);
	}
	writeBuffer(&trainData.B_data);
	
	fclose(D);

	srand(time(NULL));
	return;
}
char selectF() {
	return setF.symb[rand()%setF.Fsize];
}
char selectT() {
	return setF.symb[rand()%setF.Tsize + setF.Fsize];
}
char selectS() {
	return setF.symb[rand()%(setF.Tsize + setF.Fsize)];
}

char* genChromo(int _length, int _genes, char* _writeTo) {
	int geneLen = _length/_genes;
	int head = (geneLen - 1)/FSET_MAX_ARGS;

	int i = 0;
	for (int g = 0; g < _genes; g++) {
		_writeTo[i] = selectF();
		for (i++; i < g*geneLen + head; i++)
			_writeTo[i] = selectS();
		for (; i < (g+1)*geneLen; i++)
			_writeTo[i] = selectT();
	}

	return _writeTo + i;
}
void printChromo(int _length, char* _read) {
	char temp = *(_read+_length);
	*(_read + _length) = '\0';

	printf("%s\n", _read);
	*(_read + _length) = temp;

	return;
}
void copyChromo(char* _srcPool, int _srcChro, char* _tgtPool, int _tgtChro, int _len) {
	char* src = _srcPool + _srcChro * _len;
	char* tgt = _tgtPool + _tgtChro * _len;

	for (int i = 0; i < _len; i++)
		tgt[i] = src[i];

	return;
}
	
population initPop(int _length, int _count, int _genes) {
	population P;

	{
		int trimmed = 0;
		int rem = _length % _genes;
		if (rem) {
			_length -= rem;
			trimmed = 1;
		}
		int gLen = _length / _genes;
		if (gLen % 2 == 0) {
			_length -= _genes;
			trimmed = 1;
		}
		if (trimmed)
			printf("Warning: Chromosome parameter mismatch. Length was trimmed to %i.\n", _length);
	}

	P.length = _length;
	P.count = _count;
	P.genes = _genes;

	P.poolA = (char*) malloc ( sizeof(char) * _length * _count);
	P.poolB = (char*) malloc ( sizeof(char) * _length * _count);

	P.B_pool = 	initBuffer(sizeof(char) *_count *_length, 		INPUT_BUFFER);
	P.B_coeffs = 	initBuffer(sizeof(float) *_genes, 			OUTPUT_BUFFER);
	P.B_eval = 	initBuffer(sizeof(float) * trainData.size *_genes,	INTERNAL_BUFFER);
	P.B_assembly = 	initBuffer(sizeof(char) *_count *(_length +_genes) * 4,	INTERNAL_BUFFER);
	P.B_MLS = 	initBuffer(sizeof(float) *(_genes + 1) *_genes, 	INTERNAL_BUFFER);
	P.B_fitness = 	initBuffer(sizeof(float) *_count,		 	OUTPUT_BUFFER);

	P.fitness = P.B_fitness.data;
	switchPool(&P);

	char* writeTo = P.pool;
	for (int i = 0; i < _count; i++) {
		writeTo = genChromo(_length, _genes, writeTo);
	}

	loadRateParams(&P);

	return P;
}
void deletePop(population* _P) {
	free( (*_P).poolA );
	(*_P).poolA = NULL;

	free( (*_P).poolB );
	(*_P).poolB = NULL;

	free( (*_P).fitness );
	(*_P).fitness = NULL;

	switchPool(_P);

	freeBuffer( &(*_P).B_eval );
	freeBuffer( &(*_P).B_pool );
	freeBuffer( &(*_P).B_assembly );
	freeBuffer( &(*_P).B_coeffs );
	freeBuffer( &(*_P).B_MLS );

	return;
}
void popDump(population* _P) {
	int 	c = (*_P).count;
	int 	l = (*_P).length;
	char* 	p = (*_P).pool;

	for (int i = 0; i < c; i++) {
		for (int j = 0; j < l; j++)
			printf("%c", p[i*l+j]);
		printf("\n");
	}

	return;
}
void switchPool(population* _P) {
	if ((*_P).pool == (*_P).poolA)
		(*_P).pool = (*_P).poolB;
	else
		(*_P).pool = (*_P).poolA;

	(*_P).B_pool.data = (*_P).pool;

	return;
}
void loadRateParams(population* _P) {
	// TODO: load from file
	(*_P).mutation_r 	= 0.05;
	(*_P).transIS_r		= 0.1;
	(*_P).transRIS_r	= 0.1;
	(*_P).transG_r		= 0.1;
	(*_P).recombi1_r	= 0.2;
	(*_P).recombi2_r	= 0.5;
	(*_P).recombiG_r	= 0.1;

	return;
}
void computeFitness(population* _P) {
	static int init = 0;
	static kernel* K_compile = NULL;
	static kernel* K_parse = NULL;
	static kernel* K_buildMLS = NULL;
	static kernel* K_QR = NULL;
	static kernel* K_solve = NULL;
	static kernel* K_scale = NULL;

	if (init == 0) {
		K_compile =	find_CL("compile");
		K_parse =	find_CL("parse");
		K_buildMLS =	find_CL("build_MLS_entry");
		K_QR =		find_CL("QR_factor");
		K_solve =	find_CL("back_substitute");
		K_scale =	find_CL("get_fitness");
	}

	// COMPILE CHROMOSOMES TO INSTRUCTIONS
	writeBuffer( &(*_P).B_pool );
	if (init == 0)
		call_CL(K_compile, (*_P).count, 0, 1, "bbii", 
			&(*_P).B_pool, &(*_P).B_assembly,
			&(*_P).length, &(*_P).genes
			);
	else
		call_CL(K_compile, (*_P).count, 0, 0);
	// EVALUATE INSTRUCTIONS WITH TRAINING DATA
	for (int i = 0; i < (*_P).count; i++) {
		call_CL(K_parse, trainData.size, 0, 1, "biibbii",
			&(*_P).B_assembly, &i, &(*_P).length, &trainData.B_data,
			&(*_P).B_eval, &trainData.args, &(*_P).genes
			);
	// BUILD METHOD-OF-LEAST-SQUARES MATRIX
		for (int r = 0; r < (*_P).genes; r++)
			for (int c = 0; c <= (*_P).genes; c++)
				call_CL(K_buildMLS, trainData.size, 0, 1, "bbbiiiii",
				&(*_P).B_eval, &trainData.B_data, &(*_P).B_MLS,
				&trainData.args, &trainData.size, &(*_P).genes,
				&r, &c
				);
		if (init == 0) {
	// SOLVE FOR LEAST-SQUARES COEFFICIENTS
			call_CL(K_QR, (*_P).genes/2, 0, 1, "bi",
				&(*_P).B_MLS, &(*_P).genes
				);
			call_CL(K_solve, (*_P).genes, 0, 1, "bbi",
				&(*_P).B_MLS, &(*_P).B_coeffs, &(*_P).genes
				);
		}
		else {	
			call_CL(K_QR, (*_P).genes/2, 0, 0);
			call_CL(K_solve, (*_P).genes, 0, 0);
		}
	// SCALE RESULTS AND COMPUTE FITNESS
		call_CL(K_scale, trainData.size, 0, 1, "bbbbiiii",
			&(*_P).B_eval, &(*_P).B_coeffs, &trainData.B_data, &(*_P).B_fitness,
			&trainData.size, &(*_P).genes, &trainData.args, &i
			);
	}

	if (init == 0)
		init = 1;
	return;
}
void replicate(population* _P) {
	static int init = 0;
	static kernel* Kscan = NULL;

	if (init == 0) {
		Kscan = find_CL("scan");
		call_CL(Kscan, (*_P).count, 0, 1, "bi",
			&(*_P).B_fitness, &(*_P).count
			);
	}
	else {
		call_CL(Kscan, (*_P).count, 0, 0);
	}
	readBuffer(&(*_P).B_fitness);

	for (int c = 0; c < (*_P).count; c++) {
		float s = (float)rand()/INT_MAX * (*_P).fitness[(*_P).count-1];
		int i = 0;
		for (; i < (*_P).count; i++)
			if ( (*_P).fitness[i] >= s )
				break;
		copyChromo( (*_P).pool, i, 
				(*_P).pool == (*_P).poolA? (*_P).poolB: (*_P).poolA,
				c, (*_P).length
				);		
	}
	switchPool(_P);

	if (init == 0)
		init = 1;
	return;
}
void mutate(population* _P) {
	int targets = (int) ( (*_P).mutation_r * (*_P).count );
	int len = (*_P).length;
	int gLen = len / (*_P).genes;
	int h = (gLen - 1)/FSET_MAX_ARGS;
	
	select_noReplace( (*_P).count );
	for (int t = 0; t < targets; t++) {
		int s = select_noReplace(0);
		int c = rand() % len;
		if (c % gLen == 0)
			(*_P).pool[s*len + c] = selectF();
		else if (c % gLen < h)
			(*_P).pool[s*len + c] = selectS();
		else
			(*_P).pool[s*len + c] = selectT();
	}

	return;
}
void transpose(population* _P) {
	int l = (*_P).length;
	int g = (*_P).genes;
	int n = (*_P).count;
	char* p = (*_P).pool;

	int l_g = l/g;	// length of gene
	int h = (l_g-1)/FSET_MAX_ARGS;

	int targets;	// number of targets
	int tgt_i; 	// target chromosome
	int tgt_g; 	// target gene wrt chromosome
	int src_c; 	// source start wrt chromosome
	int tgt_c; 	// target start wrt gene
	int l_t;	// transposon length
	void shift_copy() {
		for (int i = 0; i < l_t; i++)
			scratch[i] = p[l*tgt_i + src_c + i];
		for (int i = h-1; i >= tgt_c + l_t; i--)
			p[l*tgt_i + l_g*tgt_g + i] = p[l*tgt_i + l_g*tgt_g + i - l_t];
		for (int i = 0; i < l_t; i++)
			p[l*tgt_i + l_g*tgt_g + tgt_c + i] = scratch[i];
		return;
	}

	select_noReplace(n);

	// IS
	targets = (int)( (*_P).transIS_r * n);
	for (int x = 0; x < targets; x++) {
		tgt_i = select_noReplace(0);
		tgt_g = rand() % g;
		src_c = rand() % l;
		tgt_c = rand() % (h-1) + 1;
		l_t = (rand() % 3 + 1) % (l_g - tgt_c + 1);

		shift_copy();

	}

	// RIS
	targets = (int) ( (*_P).transRIS_r * n);
	for (int x = 0; x < targets; x++) {
		tgt_i = select_noReplace(0);
		tgt_g = rand() % g;
		src_c = l_g*tgt_g + rand() % h;
		tgt_c = 0;

		int found_symb = 0;
		while ( !found_symb && src_c < h ) {
			for (int i = 0; i < setF.Fsize; i++) {
				if ( p[l*tgt_i + l_g*tgt_g + src_c] == setF.symb[i] ) {
					found_symb = 1;
					break;
				}
			}
			src_c++;
		}
		if (!found_symb)
			continue;
		src_c--;

		l_t = (rand() % 3 + 1) % (h - (src_c % l_g) + 1);

		shift_copy();
	}

	// G
	targets = (int) ( (*_P).transG_r * n);
	for (int x = 0; x < targets; x++) {
		tgt_i = select_noReplace(0);
		tgt_g = rand() % g;
		
		for (int i = 0; i < l_g; i++)
			scratch[i] = p[l*tgt_i + i];
		for (int i = 0; i < l_g; i++)
			p[l*tgt_i + i] = p[l*tgt_i + l_g*tgt_g + i];
		for (int i = 0; i < l_g; i++)
			p[l*tgt_i + l_g*tgt_g + i] = scratch[i];
	}
}
void recombine(population* _P) {
	int l = (*_P).length;
	int g = (*_P).genes;
	int n = (*_P).count;
	char* p = (*_P).pool;
	int l_g = l/g;	

	int targets;
	int L_i, R_i;
	int x1, x2;
	void cut_paste() {
		for (int i = x1; i < x2; i++) {
			scratch[i] = p[l*L_i + i];
			p[l*L_i + i] = p[l*R_i + i];
			p[l*R_i + i] = scratch[i];
		}
		return;
	}

	select_noReplace(n);

	// 1 Point
	targets = (int)( (*_P).recombi1_r * n )/2;
	for (int x = 0; x < targets; x++) {
		L_i = select_noReplace(0);
		R_i = select_noReplace(0);
		x1 = 0;
		x2 = rand() % l;
		cut_paste();
	}

	// 2 Point
	targets = (int)( (*_P).recombi2_r * n )/2;
	for (int x = 0; x < targets; x++) {
		L_i = select_noReplace(0);
		R_i = select_noReplace(0);
		int a = rand() % l;
		int b = rand() % l;
		x1 = a < b? a: b;
		x2 = a < b? b: a;
		cut_paste();
	}

	// Gene
	targets = (int)( (*_P).recombiG_r * n )/2;
	for (int x = 0; x < targets; x++) {
		L_i = select_noReplace(0);
		R_i = select_noReplace(0);
		x1 = (rand() % g) * l_g;
		x2 = x1 + l_g;
		cut_paste();
	}
	
}
void rankSurvivors(population* _P) {
	int n = (*_P).count;

	computeFitness(_P);
	readBuffer(&(*_P).B_fitness);

	int s = 0;
	for (int i = 0; i < n; i++)
		if ((*_P).fitness[i] > (*_P).fitness[s])
			s = i;
	
	//printChromo( (*_P).length, (*_P).pool + (*_P).length*s ); //TODO

	kernel* K_compile = 	find_CL("compile");
	kernel* K_parse =	find_CL("parse");
	kernel* K_buildMLS =	find_CL("build_MLS_entry");
	kernel* K_QR =		find_CL("QR_factor");
	kernel* K_solve =	find_CL("back_substitute");

	
	call_CL(K_compile, (*_P).count, 0, 1, "bbii", 
		&(*_P).B_pool, &(*_P).B_assembly,
		&(*_P).length, &(*_P).genes
		);
	call_CL(K_parse, trainData.size, 0, 1, "biibbii",
		&(*_P).B_assembly, &s, &(*_P).length, &trainData.B_data,
		&(*_P).B_eval, &trainData.args, &(*_P).genes
		);
	for (int r = 0; r < (*_P).genes; r++)
		for (int c = 0; c <= (*_P).genes; c++)
			call_CL(K_buildMLS, trainData.size, 0, 1, "bbbiiiii",
			&(*_P).B_eval, &trainData.B_data, &(*_P).B_MLS,
			&trainData.args, &trainData.size, &(*_P).genes,
			&r, &c
			);
	call_CL(K_QR, (*_P).genes/2, 0, 1, "bi",
		&(*_P).B_MLS, &(*_P).genes
		);
	call_CL(K_solve, (*_P).genes, 0, 1, "bbi",
		&(*_P).B_MLS, &(*_P).B_coeffs, &(*_P).genes
		);

	readBuffer(&(*_P).B_coeffs);

	solution Y = constructSolution( (float*)(*_P).B_coeffs.data, 
					(*_P).genes, 
					(*_P).pool + s*(*_P).length, 
					(*_P).length 
					);

	printSolution(&Y);

	return;
}

solution constructSolution(float* _coeffs, int _n_coeffs, char* _chromo, int _length) {
	solution S;

	S.size = _n_coeffs;
	S.terms = (term*) malloc (sizeof(term)*_length);
	S.coeffs = (float*) malloc (sizeof(float)*_n_coeffs);

	for (int i = 0; i < _n_coeffs; i++)
		S.coeffs[i] = _coeffs[i];

	int read_ptr = 0;
	int write_ptr = 0;
	int addr_ptr = 1;

	int gLen = _length/_n_coeffs;
	int genes_remain = S.size;
	int symbs_remain = 1;
	int n_args = 0;
	while (genes_remain--) {
		symbs_remain = 1;
		while (symbs_remain--) {
			switch (_chromo[read_ptr]) {
				case '+':
					S.terms[write_ptr].L = S.terms + addr_ptr++;
					S.terms[write_ptr].R = S.terms + addr_ptr++;
					n_args = 2;
				break;
				case '-':
					S.terms[write_ptr].L = S.terms + addr_ptr++;
					S.terms[write_ptr].R = S.terms + addr_ptr++;
					n_args = 2;
				break;
				case '*':
					S.terms[write_ptr].L = S.terms + addr_ptr++;
					S.terms[write_ptr].R = S.terms + addr_ptr++;
					n_args = 2;
				break;
				case '/':
					S.terms[write_ptr].L = S.terms + addr_ptr++;
					S.terms[write_ptr].R = S.terms + addr_ptr++;
					n_args = 2;
				break;
				case 'Q':
					S.terms[write_ptr].L = S.terms + addr_ptr++;
					S.terms[write_ptr].R = NULL;
					n_args = 1;
				break;
				case 'x':
					S.terms[write_ptr].L = NULL;
					S.terms[write_ptr].R = NULL;
					n_args = 0;
				break;
				default:
					S.terms[write_ptr].L = NULL;
					S.terms[write_ptr].R = NULL;
					n_args = 0;
				break;
			}
			S.terms[write_ptr].expr = _chromo[read_ptr];
			symbs_remain += n_args;
			read_ptr += 1;
			write_ptr += 1;
		}
		read_ptr = gLen * (_n_coeffs - genes_remain);
		addr_ptr++;
	}

	return S;
}
void printSolution(solution* _S) {

	int printTerm(term* _T) {
		static int last_ptr = 0;
		if (_T == NULL) {
			last_ptr = 0;
			return 0;
		}
		switch ((*_T).expr) {
			case '+':
				printf("(");
				printTerm( (*_T).L );
				printf("+");
				printTerm( (*_T).R );
				printf(")");
			break;
			case '-':
				printf("(");
				printTerm( (*_T).L );
				printf("-");
				printTerm( (*_T).R );
				printf(")");
			break;
			case '*':
				printf("(");
				printTerm( (*_T).L );
				printf("*");
				printTerm( (*_T).R );
				printf(")");
			break;
			case '/':
				printf("(");
				printTerm( (*_T).L );
				printf("/");
				printTerm( (*_T).R );
				printf(")");
			break;
			case 'Q':
				printf("sqrt(");
				printTerm( (*_T).L );
				printf(")");
			break;
			case 'x':
				printf("x");
			break;
		}
		last_ptr++;
		return last_ptr;
	}

	int last_ptr = 0;
	printTerm(NULL);
	for (int i = 0; i < (*_S).size; i++) {
		printf("(%2.2f)*", (*_S).coeffs[i]);
		last_ptr = printTerm( (*_S).terms + last_ptr );
		if (i < (*_S).size-1)
			printf(" + ");
	}

	return;
}

int select_noReplace(int _size) {
	static int selectionRange[4096];
	static int selectionMax = 0;

	int ret = -1;
	if (_size) {
		for (int i = 0; i < _size; i++)
			selectionRange[i] = i;
		selectionMax = _size;
	}
	else {
		if (selectionMax == 0)
			return -1;
		int s = rand() % selectionMax;
		ret = selectionRange[s];
		selectionRange[s] = selectionRange[selectionMax-1];
		selectionMax--;
	}

	return ret;
}
